import cv2
import numpy as np
from detection import detect_on_windows, nms_boxes
from visualization import draw_boxes
from ultralytics import YOLO

def main():
    image_path = 'your_big_image.jpg'  # 你的大图路径
    model_path = 'yolov8n.pt'           # YOLO模型路径
    window_size = 640                    # 窗口大小
    step = 320                           # 步长
    conf_thres = 0.25                   # 置信度阈值
    iou_thres = 0.5                     # IOU阈值

    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        print(f"Error: Unable to read image at {image_path}")
        return

    ori_image = image.copy()
    
    # 1. 计算填充大小
    padding = window_size // 2  # 可以选择窗口的一半作为填充大小
    padded_image = cv2.copyMakeBorder(image, padding, padding, padding, padding, cv2.BORDER_CONSTANT, value=[0, 0, 0])

    # 2. 检测窗口
    all_boxes = detect_on_windows(padded_image, window_size, step, model_path, conf_thres)

    # 3. 映射检测框坐标回原图
    for box in all_boxes:
        box[0] -= padding  # x1
        box[1] -= padding  # y1
        box[2] -= padding  # x2
        box[3] -= padding  # y2

    nmsed_boxes = nms_boxes(all_boxes, iou_threshold=iou_thres)

    # 加载模型并获取类别名称
    model = YOLO(model_path)
    class_names = model.names if hasattr(model, 'names') else None

    # 绘制检测框
    result_img = draw_boxes(ori_image, nmsed_boxes, class_names=class_names)

    # 保存和显示结果
    cv2.imwrite('result.jpg', result_img)
    cv2.imshow('Detection Result', result_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()

